

import numpy as np
import os
import pickle


def gen_list_5fold(candi_dir,
                   candis_thresh,
                   pkl_dir):
    candidates_list = []
    for fold_idx in range(1, 6):
        fold_name = 'fold' + str(fold_idx)
        candi_fold_dir = os.path.join(candi_dir, fold_name)
        scan_names = os.listdir(candi_fold_dir)
        for scan_name in scan_names:
            if scan_name == 'b8bb02d229361a623a4dc57aa0e5c485':
                continue
            try:
                candis = np.load(os.path.join(candi_fold_dir, scan_name, 'candis.npy'))
            except:
                continue
            for icandi in xrange(len(candis)):  # for every slice
                candi_slice = candis[icandi]
                if len(candi_slice) == 0:
                    continue

                for bbox in candi_slice:  # for every candidate in one slice
                    if bbox[4] < candis_thresh:
                        continue
                    if scan_name[0] == 'L':
                        source_str = 'lidc'
                    else:
                        source_str = 'kaggle_mali'
                    bb = [np.floor(bbox[1]), np.floor(bbox[0]), icandi, np.ceil(bbox[3])+1, np.ceil(bbox[2])+1, icandi+1]
                    bb2 = np.array(bb, dtype=np.int)
                    candidate_dict = dict(scan_id=scan_name, source=source_str, bbox=bb2)
                    candidates_list.append(candidate_dict)
    with open(os.path.join(pkl_dir, 'candis_v9_5fold.pkl'), 'wb') as f:
        pickle.dump(candidates_list, f, -1)


def gen_list_one_fold(candi_dir,
                      fold_name,
                      candis_thresh,
                      pkl_dir):
    candidates_list = []

    candi_fold_dir = os.path.join(candi_dir, fold_name)
    scan_names = os.listdir(candi_fold_dir)
    for scan_name in scan_names:
        if scan_name == 'b8bb02d229361a623a4dc57aa0e5c485':
            continue
        try:
            candis = np.load(os.path.join(candi_fold_dir, scan_name, 'candis.npy'))
        except:
            continue
        for icandi in xrange(len(candis)):  # for every slice
            candi_slice = candis[icandi]
            if len(candi_slice) == 0:
                continue

            for bbox in candi_slice:  # for every candidate in one slice
                if bbox[4] < candis_thresh:
                    continue
                source_str = fold_name
                bb = [np.floor(bbox[1]), np.floor(bbox[0]), icandi, np.ceil(bbox[3])+1, np.ceil(bbox[2])+1, icandi+1]
                bb2 = np.array(bb, dtype=np.int)
                candidate_dict = dict(scan_id=scan_name, source=source_str, bbox=bb2)
                candidates_list.append(candidate_dict)
    with open(os.path.join(pkl_dir, 'candis_v9_' + fold_name+'.pkl'), 'wb') as f:
        pickle.dump(candidates_list, f, -1)

if __name__ == '__main__':
    # gen_list_onefold()
    pass